# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os
import math
import datetime

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from coin.base.datetime_util import to_datetime
from coin.support.feed_tool.feed_stats.logic.util import datetime_span
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct

FLAGS = flags.FLAGS


def plot_price_vs_limit(df, csv_root, mid_price, limit_low, limit_high, trading_date_str, symbol):
  plt.plot(df.index,
           df[mid_price],
           'kx',
           df.index,
           df[limit_low],
           'r-',
           df.index,
           df[limit_high],
           'b-',
           lw=0.2,
           markersize=0.2)

  plt.title('Okex Futures %s %s mid_price VS limit_price' % (trading_date_str, symbol))
  plt.grid(True)
  plt.savefig(
      os.path.join(csv_root,
                   'OkexFutures_%s_%s_midprice_vs_limitprice.png' % (trading_date_str, symbol)))
  plt.close()


def simple_stats_and_print(df, mid_price, limit_low, limit_high, trading_date_str, symbol):
  percentiles = [.005, .1]
  print(symbol)
  print(trading_date_str)
  df['mid_low_spread'] = (df[mid_price] - df[limit_low]) / df[mid_price]
  df['mid_high_spread'] = (df[limit_high] - df[mid_price]) / df[mid_price]
  df['low_high_spread'] = (df[limit_high] - df[limit_low]) / df[mid_price]
  print(df['mid_low_spread'].describe(percentiles=percentiles))
  print(df['mid_high_spread'].describe(percentiles=percentiles))
  print(df['low_high_spread'].describe(percentiles=percentiles))


def plot_col1_vs_col2(df, csv_root, col1, col2, trading_date_str, query_machine, symbol):
  fig, ax1 = plt.subplots(figsize=(20, 10))

  ax2 = ax1.twinx()
  plt.title('Okex Futures %s %s %s %s VS %s' %
            (trading_date_str, query_machine, symbol, col1, col2))
  ax1.plot(df.index, df[col1], color='g', marker='.')
  ax2.plot(df.index, df[col2], color='r', marker='.')
  tick_space = (ax1.get_yticks()[-1] - ax1.get_yticks()[0]) / 20
  tick_space = int(math.ceil(tick_space / 5)) * 5
  ax1.yaxis.set_major_locator(ticker.MultipleLocator(tick_space))

  ax1.set_ylabel(col1, color='g')
  ax1.grid(True)
  ax2.set_ylabel(col2, color='r')
  plt.savefig(
      os.path.join(
          csv_root,
          'OkexFutures_%s_%s_%s_%s_vs_%s.png' %
          (trading_date_str, query_machine, symbol, col1, col2)))
  plt.close()


def plot_df(df, csv_root, trading_date_str, query_machine, symbol):
  plot_col1_vs_col2(df,
                    csv_root,
                    'mid_price',
                    'open_interest',
                    trading_date_str,
                    query_machine,
                    symbol)
  plot_col1_vs_col2(df, csv_root, 'mid_price', 'vol_24h', trading_date_str, query_machine, symbol)
  simple_stats_and_print(df, 'mid_price', 'limit_low', 'limit_high', trading_date_str, symbol)
  plot_price_vs_limit(df,
                      csv_root,
                      'mid_price',
                      'limit_low',
                      'limit_high',
                      trading_date_str,
                      symbol)
  plot_col1_vs_col2(df,
                    csv_root,
                    'mid_price',
                    'long_short_position_account_ratio',
                    trading_date_str,
                    query_machine,
                    symbol)
  plot_col1_vs_col2(df,
                    csv_root,
                    'mid_price',
                    'long_short_margin_ratio',
                    trading_date_str,
                    query_machine,
                    symbol)


def hour_describe(df, trading_date_str, query_machine, symbol):
  pandas.set_option('display.width', 1000)
  start = to_datetime(int(df.iloc[1]['timestamp_x']))
  end = to_datetime(int(df.iloc[-1]['timestamp_x']))
  prev_dt = start
  start_hour = start.hour
  start_hour = start.replace(hour=start_hour + 1, minute=0, second=0, microsecond=0)
  end_hour = end.replace(minute=0, second=0, microsecond=0)
  print('Hour describe for %s %s %s:' % (trading_date_str, query_machine, symbol))
  # timestamp column does not need to describe.
  df = df.drop('timestamp_x', axis=1)
  for dt in datetime_span(start_hour, end_hour, datetime.timedelta(hours=1)):
    print('%s - %s' % (prev_dt, dt))
    print(df.between_time(prev_dt.time(), dt.time()).describe())
    print()
    prev_dt = dt

  # For the last hour.
  print('%s - %s' % (prev_dt, end))
  print(df.between_time(prev_dt.time(), end.time()).describe())
  print()


def daily_describe(df, trading_date_str, query_machine, symbol):
  pandas.set_option('display.width', 1000)
  start = to_datetime(int(df.iloc[1]['timestamp_x']))
  end = to_datetime(int(df.iloc[-1]['timestamp_x']))
  print('Daily describe for %s %s %s:' % (trading_date_str, query_machine, symbol))
  df = df.drop('timestamp_x', axis=1)
  print('%s - %s' % (start, end))
  print(df.between_time(start.time(), end.time()).describe())
  print()


def process_dfs(book_df, ticker_df, account_ratio_df, margin_ratio_df):
  book_df = sample_df_per_second(book_df)
  ticker_df = sample_df_per_second(ticker_df)
  account_ratio_df = sample_df_per_second(account_ratio_df)
  margin_ratio_df = sample_df_per_second(margin_ratio_df)

  book_df['mid_price'] = (book_df['bid0'] + book_df['ask0']) / 2
  df = pandas.merge(book_df, ticker_df, how='inner', left_index=True, right_index=True)
  df = pandas.merge(df, account_ratio_df, how='inner', left_index=True, right_index=True)
  df = pandas.merge(df, margin_ratio_df, how='inner', left_index=True, right_index=True)
  return df


def sample_df_per_second(df):
  return df.resample('S', closed='left', label='left').ffill()


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  assert len(df) > 0, 'Empty instrument csv: %s' % csv_dir
  return df


def read_csv_into_df2(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['fetch_time'], format='%Y-%m-%d %H:%M:%S.%f')
  df = df.set_index('datetime')
  assert len(df) > 0, 'Empty instrument csv: %s' % csv_dir
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  trading_date_str = FLAGS.trading_date
  trading_date = datetime.datetime.strptime(trading_date_str, '%Y%m%d')
  query_machine = FLAGS.query_machine
  symbols = flags.FLAGS.symbols.split(",")

  for symbol in symbols:
    product = OkexFuturesProduct.FromStr(symbol, trading_date)
    currency = product.base.symbol
    csv_book = 'book/%s.csv' % product.full_symbol
    csv_ticker = 'ticker_%s_%s.csv' % (product.symbol, trading_date_str)
    csv_long_short_position_account_ratio = \
        'long_short_position_account_ratio_%s_%s_%s.csv' % \
        (trading_date_str, query_machine, currency)
    csv_long_short_margin_ratio = \
        'long_short_margin_ratio_%s_%s_%s.csv' % \
        (trading_date_str, query_machine, currency)

    book_df = read_csv_into_df(csv_root, csv_book)
    ticker_df = read_csv_into_df(csv_root, csv_ticker)
    account_ratio_df = read_csv_into_df2(csv_root, csv_long_short_position_account_ratio)
    margin_ratio_df = read_csv_into_df2(csv_root, csv_long_short_margin_ratio)

    df = process_dfs(book_df, ticker_df, account_ratio_df, margin_ratio_df)

    # Put output plot into same csv_root directory.
    plot_df(df, csv_root, trading_date_str, query_machine, currency)
    daily_describe(df, trading_date_str, query_machine, currency)
    hour_describe(df, trading_date_str, query_machine, currency)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv root directory.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('query_machine', None, 'Instance machine name.')

  flags.DEFINE_string('symbols', 'BTC-USD.QUARTER,ETH-USD.QUARTER,EOS-USD.QUARTER', 'symbol name')

  app.run(main)
